package edu.hit.irlab.coref;

import java.io.IOException;
import java.util.ArrayList;

import nlpeap.ml.OfflineClassifier;
import nlpeap.ml.weka.WEKAInstanceClassifier;
import nlpeap.tms.DefaultDocumentfactory;
import nlpeap.tms.DocumentFactory;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dom4j.Document;
import org.dom4j.DocumentException;
import org.dom4j.Element;

import edu.hit.irlab.coref.corpus.Corpus;
import edu.hit.irlab.coref.mentionfinder.DefaultCorpus_Mentionfinder;
import edu.hit.irlab.coref.mentionfinder.Mention;
import edu.hit.irlab.coref.mentionfinder.MentionFinder;
import edu.hit.irlab.coref.postprocess.Postprocessor;
import edu.hit.irlab.coref.resolve.BinaryClassifyResolver;
import edu.hit.irlab.coref.resolve.Entity;
import edu.hit.irlab.coref.resolve.Resolver;
import edu.hit.irlab.coref.resolve.algorithm.CloseFirst;
import edu.hit.irlab.coref.resolve.algorithm.ResolveAlgorithm;
import edu.hit.irlab.coref.visualize.MapDoctodom;

public class TestModel {
	public static Log _log = LogFactory.getLog(TestModel.class.getName());
	private static final CRConfigProperties _cprop = CRConfigProperties.getInstance();
	private static final Postprocessor _postProcessor = _cprop.getPostPrecessor();
	
	private MentionFinder<Document> _mentionFinder; 
	private ResolveAlgorithm _resolveAlgorithm;
	
	public TestModel(
			MentionFinder<Document> mentionFinder, 
			ResolveAlgorithm resolveAlgorithm
		) {
		this._mentionFinder = mentionFinder;
		this._resolveAlgorithm = resolveAlgorithm;
	}
	
	public void addtodom(ArrayList<Entity> entitychains,
			 org.dom4j.Document domobj) {
		Element rootnode = domobj.getRootElement();
		
		if (entitychains.size() < 1)
			return;
		Element corefnode = rootnode.addElement("coref");
		int crid = 0;
		for (Entity entity : entitychains) {
			Element crnode = corefnode.addElement("cr");
			crnode.addAttribute("id", String.valueOf(crid++));
			int mentionid = 0;
			for (Mention mention : entity.getmentions()) {
				Element mentionnodea = crnode.addElement("mention");
				mentionnodea.addAttribute("id", String.valueOf(mentionid++));
				mentionnodea.addAttribute("beg", 
						String.valueOf(mention.getExtentBegintokenId()));
				mentionnodea.addAttribute("end", 
						String.valueOf(mention.getExtentEndtokenId()));
				mentionnodea.addAttribute("type", 
						mention.getMentionSubType());
			}
		}
	}
	
	public Document testCorpus(
			String corpusPath) throws IOException, DocumentException {
		Document dom = Corpus.loadCorpus(corpusPath);
		Mention[] mentions = _mentionFinder.findmentions(dom);
		ArrayList<Entity> entityRes = _resolveAlgorithm.resolve(mentions);
		entityRes = _postProcessor.postProcess(entityRes);
		
		nlpeap.tms.Document doc;
		if (mentions.length <= 0) {
			//if there is no Mention in the Corpus
			String text_cont = Corpus.getTextCont(dom);
			DocumentFactory<String> docFac = new DefaultDocumentfactory(null);
			doc = docFac.createDocument(text_cont);
		} else
			doc = mentions[0].getdoc();
		
		Document visDom = MapDoctodom.mapdocumentobjtodom(doc);
		addtodom(entityRes, visDom);
		
		return visDom;
	}
	
	public static void main(String[] argv) throws 
		ClassNotFoundException, InstantiationException, 
		IllegalAccessException, IOException, DocumentException {
		MentionFinder<Document> corpusMentionFinder = new DefaultCorpus_Mentionfinder();
		CRConfigProperties prop = CRConfigProperties.getInstance();
		ArrayList<PairFeatureExtractor> fes = prop.getPairFEs();
		OfflineClassifier model = new WEKAInstanceClassifier(prop.getModelName());
		Resolver[] resolvers = new Resolver[]{
				new BinaryClassifyResolver(fes, model)
			};
		ResolveAlgorithm resolveAlgorithm = new CloseFirst(resolvers);
		TestModel ins = new TestModel(corpusMentionFinder, resolveAlgorithm);
		
		String corpusPath = "e:/ll/a32.xml";
		System.out.println(ins.testCorpus(corpusPath).asXML());
	}
}
